﻿from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Dict, Any, Optional, List
import pandas as pd
import numpy as np
import joblib
import os
import uvicorn
from PIL import Image
from io import BytesIO
import tempfile
from tensorflow.keras.models import load_model
import shutil

app = FastAPI(title="集成医疗诊断平台API")

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 定义请求模型
class BreastCancerRequest(BaseModel):
    radius_mean: float
    texture_mean: float
    perimeter_mean: float
    area_mean: float
    smoothness_mean: float
    compactness_mean: float
    concavity_mean: float
    concave_points_mean: float
    symmetry_mean: float
    radius_se: float
    perimeter_se: float
    area_se: float
    compactness_se: float
    concavity_se: float
    concave_points_se: float
    fractal_dimension_se: float
    radius_worst: float
    texture_worst: float
    perimeter_worst: float
    area_worst: float
    smoothness_worst: float
    compactness_worst: float
    concavity_worst: float
    concave_points_worst: float
    symmetry_worst: float
    fractal_dimension_worst: float

class CervicalCancerRequest(BaseModel):
    Age: int
    Number_of_sexual_partners: int
    First_sexual_intercourse: int
    Smokes_years: int
    Hormonal_Contraceptives_years: float
    IUD: int
    IUD_years: float
    STDs: int
    STDs_number: int
    STDs_genital_herpes: int
    STDs_HIV: int
    STDs_condylomatosis: int
    STDs_vaginal_condylomatosis: int
    STDs_vulvo_perineal_condylomatosis: int
    STDs_AIDS: int
    Dx_Cancer: int
    Dx_CIN: int
    Hinselmann: int
    Schiller: int
    Citology: int

class PredictionResponse(BaseModel):
    label: int
    probability: float
    
# 乳腺癌预测函数
def predict_breast_cancer(data: Dict[str, float]):
    # 加载模型
    model_path = os.path.join('model', 'mlp_breast_cancer_model.pkl')
    model = joblib.load(model_path)
    
    # 将数据转换为DataFrame并保持列的顺序
    columns = [
        'radius_mean','texture_mean','perimeter_mean','area_mean','smoothness_mean',
        'compactness_mean','concavity_mean','concave points_mean','symmetry_mean',
        'radius_se','perimeter_se','area_se','compactness_se','concavity_se','concave points_se',
        'fractal_dimension_se','radius_worst','texture_worst','perimeter_worst','area_worst',
        'smoothness_worst','compactness_worst','concavity_worst','concave points_worst',
        'symmetry_worst','fractal_dimension_worst'
    ]
    
    # 调整键名以匹配模型期望的输入
    data_adjusted = {k.replace('concave_points', 'concave points'): v for k, v in data.items()}
    df = pd.DataFrame([data_adjusted], columns=columns)
    
    # 预测
    pred_label = model.predict(df)[0]
    pred_proba = model.predict_proba(df)[0]
    
    return {
        "label": int(pred_label),
        "probability": float(pred_proba[pred_label])
    }

# 宫颈癌预测函数
def predict_cervical_cancer(data: Dict[str, Any]):
    # 加载模型
    model_path = os.path.join('model', 'best_xgb_pipeline_model.joblib')
    model = joblib.load(model_path)
    
    # 将数据转换为DataFrame并保持列的顺序
    columns = [
        'Hormonal Contraceptives (years)', 'Dx:Cancer', 'STDs (number)', 'STDs:genital herpes',
        'STDs:HIV', 'STDs:condylomatosis', 'STDs', 'Hinselmann', 'Dx:CIN', 'Citology',
        'IUD', 'IUD (years)', 'STDs:AIDS', 'Number of sexual partners', 'Age',
        'STDs:vaginal condylomatosis', 'STDs:vulvo-perineal condylomatosis',
        'Smokes (years)', 'First sexual intercourse', 'Schiller'
    ]
    
    # 将前端字段映射到模型期望的字段
    mapping = {
        'Age': 'Age',
        'Number_of_sexual_partners': 'Number of sexual partners',
        'First_sexual_intercourse': 'First sexual intercourse',
        'Smokes_years': 'Smokes (years)',
        'Hormonal_Contraceptives_years': 'Hormonal Contraceptives (years)',
        'IUD': 'IUD',
        'IUD_years': 'IUD (years)',
        'STDs': 'STDs',
        'STDs_number': 'STDs (number)',
        'STDs_genital_herpes': 'STDs:genital herpes',
        'STDs_HIV': 'STDs:HIV',
        'STDs_condylomatosis': 'STDs:condylomatosis',
        'STDs_vaginal_condylomatosis': 'STDs:vaginal condylomatosis',
        'STDs_vulvo_perineal_condylomatosis': 'STDs:vulvo-perineal condylomatosis',
        'STDs_AIDS': 'STDs:AIDS',
        'Dx_Cancer': 'Dx:Cancer',
        'Dx_CIN': 'Dx:CIN',
        'Hinselmann': 'Hinselmann',
        'Schiller': 'Schiller',
        'Citology': 'Citology'
    }
    
    # 转换数据
    model_input = {mapping[k]: v for k, v in data.items()}
    df = pd.DataFrame([model_input])
    
    # 确保列顺序与模型期望的一致
    df = df.reindex(columns=columns)
        
    # 预测
    pred_label = model.predict(df)[0]
    pred_proba = model.predict_proba(df)[0]
        
    return {
        "label": int(pred_label),
        "probability": float(pred_proba[pred_label])
    }

# 乳腺癌图片预测函数
def predict_breast_cancer_image(image_path: str):
    # 图片预处理
    def preprocess_image(img_path, target_size=(224,224)):
        try:
            img = Image.open(img_path).convert('RGB')
            img = img.resize(target_size)
            img_array = np.array(img) / 255.0
            img_array = np.expand_dims(img_array, axis=0)
            return img_array
        finally:
            # 确保图片对象被关闭
            if 'img' in locals() and hasattr(img, 'close'):
                img.close()
    
    # 加载模型
    model_path = os.path.join('model', 'best_densenet121_attention.h5')
    model = load_model(model_path, compile=False)
    
    # 类别字典
    index_to_class = {0: 'benign', 1: 'malignant', 2: 'normal'}
    
    # 图片预处理
    img_array = preprocess_image(image_path)
    
    # 预测
    pred = model.predict(img_array)
    predicted_class_index = np.argmax(pred, axis=1)[0]
    
    # 将结果转换为二分类 (0 = 良性, 1 = 恶性)
    label = 1 if predicted_class_index == 1 else 0
    probability = float(pred[0][predicted_class_index])
    
    # 清理TensorFlow资源
    import tensorflow as tf
    tf.keras.backend.clear_session()
    
    return {
        "label": label,
        "probability": probability
    }

# 宫颈癌图片预测函数
def predict_cervical_cancer_image(image_path: str):
    # 图片预处理
    def preprocess_image(img_path, target_size=(224,224)):
        try:
            img = Image.open(img_path).convert('RGB')
            img = img.resize(target_size)
            img_array = np.array(img) / 255.0
            img_array = np.expand_dims(img_array, axis=0)
            return img_array
        finally:
            # 确保图片对象被关闭
            if 'img' in locals() and hasattr(img, 'close'):
                img.close()
    
    # 加载模型
    model_path = os.path.join('model', 'best_vgg16_finetune.h5')
    model = load_model(model_path)
    
    # 类别字典
    index_to_class = {
        0: "im_Dyskeratotic",
        1: "im_Koilocytotic",
        2: "im_Metaplastic",
        3: "im_Parabasal",
        4: "im_Superficial-Intermediate"
    }
    
    # 图片预处理
    img_array = preprocess_image(image_path)
    
    # 预测
    pred = model.predict(img_array)
    predicted_class_index = np.argmax(pred, axis=1)[0]
    
    # 转换为二分类 (0 = 低风险, 1 = 高风险)
    # 根据医学判断将不同类型归为低风险或高风险
    high_risk_classes = [0, 1, 3]  # 假设这些类别是高风险
    label = 1 if predicted_class_index in high_risk_classes else 0
    probability = float(pred[0][predicted_class_index])
    
    # 清理TensorFlow资源
    import tensorflow as tf
    tf.keras.backend.clear_session()
    
    return {
        "label": label,
        "probability": probability
    }

# Excel处理函数
def process_breast_cancer_excel(file_path: str):
    # 读取Excel文件
    try:
        df = pd.read_excel(file_path) if file_path.endswith(('.xlsx', '.xls')) else pd.read_csv(file_path)
        
        # 确保包含所有必要列
        required_columns = [
            'radius_mean', 'texture_mean', 'perimeter_mean', 'area_mean', 'smoothness_mean',
            'compactness_mean', 'concavity_mean', 'concave_points_mean', 'symmetry_mean',
            'radius_se', 'perimeter_se', 'area_se', 'compactness_se', 'concavity_se', 'concave_points_se',
            'fractal_dimension_se', 'radius_worst', 'texture_worst', 'perimeter_worst', 'area_worst',
            'smoothness_worst', 'compactness_worst', 'concavity_worst', 'concave_points_worst',
            'symmetry_worst', 'fractal_dimension_worst'
        ]
        
        missing_cols = [col for col in required_columns if col not in df.columns]
        if missing_cols:
            return {"error": f"Excel文件缺少以下必要列: {', '.join(missing_cols)}"}
        
        # 使用第一行数据
        row_data = df.iloc[0].to_dict()
        
        # 调用预测函数
        return predict_breast_cancer(row_data)
    except Exception as e:
        return {"error": f"处理Excel文件时出错: {str(e)}"}

def process_cervical_cancer_excel(file_path: str):
    # 读取Excel文件
    try:
        df = pd.read_excel(file_path) if file_path.endswith(('.xlsx', '.xls')) else pd.read_csv(file_path)
        
        # 确保包含所有必要列
        required_columns = [
            'Age', 'Number_of_sexual_partners', 'First_sexual_intercourse', 'Smokes_years',
            'Hormonal_Contraceptives_years', 'IUD', 'IUD_years', 'STDs', 'STDs_number',
            'STDs_genital_herpes', 'STDs_HIV', 'STDs_condylomatosis',
            'STDs_vaginal_condylomatosis', 'STDs_vulvo_perineal_condylomatosis',
            'STDs_AIDS', 'Dx_Cancer', 'Dx_CIN', 'Hinselmann', 'Schiller', 'Citology'
        ]
        
        missing_cols = [col for col in required_columns if col not in df.columns]
        if missing_cols:
            return {"error": f"Excel文件缺少以下必要列: {', '.join(missing_cols)}"}
        
        # 使用第一行数据
        row_data = df.iloc[0].to_dict()
        
        # 调用预测函数
        return predict_cervical_cancer(row_data)
    except Exception as e:
        return {"error": f"处理Excel文件时出错: {str(e)}"}

# API路由

@app.get("/")
def read_root():
    return {"message": "欢迎使用集成医疗诊断平台API"}

# 乳腺癌诊断路由
@app.post("/api/predict/breast-cancer", response_model=PredictionResponse)
async def api_predict_breast_cancer(request: BreastCancerRequest):
    data = request.dict()
    result = predict_breast_cancer(data)
    return result

@app.post("/api/predict/breast-cancer/excel", response_model=PredictionResponse)
async def api_predict_breast_cancer_excel(file: UploadFile = File(...)):
    # 保存上传的文件
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1])
    try:
        contents = await file.read()
        with open(temp_file.name, 'wb') as f:
            f.write(contents)
        
        # 处理Excel并预测
        result = process_breast_cancer_excel(temp_file.name)
        if "error" in result:
            raise HTTPException(status_code=400, detail=result["error"])
            
        return result
    finally:
        # 安全删除临时文件
        try:
            os.unlink(temp_file.name)
        except Exception as e:
            # 忽略文件删除错误
            print(f"警告: 无法删除临时文件 {temp_file.name}: {str(e)}")

@app.post("/api/predict/breast-cancer/image", response_model=PredictionResponse)
async def api_predict_breast_cancer_image(image: UploadFile = File(...)):
    # 保存上传的图片
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(image.filename)[1])
    try:
        contents = await image.read()
        with open(temp_file.name, 'wb') as f:
            f.write(contents)
        
        # 处理图片并预测
        result = predict_breast_cancer_image(temp_file.name)
        return result
    finally:
        # 安全删除临时文件
        try:
            os.unlink(temp_file.name)
        except Exception as e:
            # 忽略文件删除错误
            print(f"警告: 无法删除临时文件 {temp_file.name}: {str(e)}")

# 宫颈癌风险评估路由
@app.post("/api/predict/cervical-cancer", response_model=PredictionResponse)
async def api_predict_cervical_cancer(request: CervicalCancerRequest):
    data = request.dict()
    result = predict_cervical_cancer(data)
    return result

@app.post("/api/predict/cervical-cancer/excel", response_model=PredictionResponse)
async def api_predict_cervical_cancer_excel(file: UploadFile = File(...)):
    # 保存上传的文件
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1])
    try:
        contents = await file.read()
        with open(temp_file.name, 'wb') as f:
            f.write(contents)
        
        # 处理Excel并预测
        result = process_cervical_cancer_excel(temp_file.name)
        if "error" in result:
            raise HTTPException(status_code=400, detail=result["error"])
            
        return result
    finally:
        # 安全删除临时文件
        try:
            os.unlink(temp_file.name)
        except Exception as e:
            # 忽略文件删除错误
            print(f"警告: 无法删除临时文件 {temp_file.name}: {str(e)}")

@app.post("/api/predict/cervical-cancer/image", response_model=PredictionResponse)
async def api_predict_cervical_cancer_image(image: UploadFile = File(...)):
    # 保存上传的图片
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(image.filename)[1])
    try:
        contents = await image.read()
        with open(temp_file.name, 'wb') as f:
            f.write(contents)
        
        # 处理图片并预测
        result = predict_cervical_cancer_image(temp_file.name)
        return result
    finally:
        # 安全删除临时文件
        try:
            os.unlink(temp_file.name)
        except Exception as e:
            # 忽略文件删除错误
            print(f"警告: 无法删除临时文件 {temp_file.name}: {str(e)}")

# 主程序
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)
